import numpy as np
import scipy as sp
import conv_matrix

"""
-   The scipy.sparse.spmatrix class might be deprecated soon, check it before running the code.
-   The exp(-jkz) convention is used. Use negative imaginary part for lossy materials.
"""


def get_reciprocal(lattice_vector_1, lattice_vector_2):
    """
    Calculate reciprocal vectors from lattice

    Args:
        lattice_vector_1 (numpy.ndarray): 1x2 vector, lattice vector 1
        lattice_vector_2 (numpy.ndarray): 1x2 vector, lattice vector 2

    Returns:
        numpy.ndarray: 2x2 matrix [b1, b2]^T, reciprocal vectors in [0] and [1] as 1x2 vectors
    """
    return 2 * np.pi * np.linalg.inv(np.array([lattice_vector_1, lattice_vector_2])).transpose()


def get_k_matrix(kx_inc, ky_inc, k1_order, k2_order, reciprocal_vector_1, reciprocal_vector_2, frequency,
                 c_const=2.998e5):
    """
    Calculate Kx, Ky matrix

    Args:
        kx_inc (float):
        ky_inc (float):
        k1_order (int):
        k2_order (int):
        reciprocal_vector_1 (numpy.ndarray):
        reciprocal_vector_2 (numpy.ndarray):
        frequency (float):
        c_const (float):

    Returns:
        tuple[scipy.sparse.spmatrix, scipy.sparse.spmatrix] : sparse matrices, including the kx and ky matrix

    """

    n_kx = 2 * k1_order + 1
    n_ky = 2 * k2_order + 1
    k0_func = 2 * np.pi * frequency / c_const
    kx_inc_norm = kx_inc / k0_func
    ky_inc_norm = ky_inc / k0_func
    n_basis = n_kx * n_ky
    kx_list = np.zeros(n_basis)
    ky_list = np.zeros(n_basis)

    for i_func in range(n_basis):
        k1 = reciprocal_vector_1 * (i_func // n_ky - k1_order)
        k2 = reciprocal_vector_2 * (i_func % n_ky - k2_order)
        kx_list[i_func] = kx_inc_norm - (k1[0] + k2[0]) / k0_func
        ky_list[i_func] = ky_inc_norm - (k1[1] + k2[1]) / k0_func

    kx_matrix = sp.sparse.diags(kx_list)
    ky_matrix = sp.sparse.diags(ky_list)

    return kx_matrix, ky_matrix


def get_kz_matrix(kx_matrix, ky_matrix, eps_inc=1, mu_inc=1, eps_trn=1, mu_trn=1):
    """
    Calculate Kz matrix

    Args:
        eps_inc (complex):
        mu_inc (complex):
        eps_trn (complex):
        mu_trn (complex):
        kx_matrix (scipy.sparse.spmatrix):
        ky_matrix (scipy.sparse.spmatrix):

    Returns:
        tuple[scipy.sparse.spmatrix, scipy.sparse.spmatrix]:
            sparse matrices, including the kz matrices in reflection and transmission areas
    """
    kx2 = kx_matrix ** 2
    ky2 = ky_matrix ** 2
    k02_ref = np.conjugate(sp.sparse.diags(np.ones(kx_matrix.shape[0]) * eps_inc * mu_inc))
    k02_trn = np.conjugate(sp.sparse.diags(np.ones(kx_matrix.shape[0]) * eps_trn * mu_trn))
    kz_inc_matrix = np.conjugate(np.sqrt((k02_ref - kx2 - ky2).astype(complex)))
    kz_trn_matrix = np.conjugate(np.sqrt((k02_trn - kx2 - ky2).astype(complex)))

    return kz_inc_matrix, kz_trn_matrix


def get_p_q_matrix(kx_matrix, ky_matrix, eps_conv_matrix, mu_conv_matrix):
    """
    Args:
        kx_matrix (scipy.sparse.spmatrix): kx matrix, diagonal
        ky_matrix (scipy.sparse.spmatrix): ky matrix, diagonal
        eps_conv_matrix (numpy.ndarray): dielectric constant convolution matrix
        mu_conv_matrix (numpy.ndarray): permeability constant convolution matrix

    Returns:
        tuple[numpy.ndarray, numpy.ndarray]: P and Q matrices
    """
    eps_conv_inv = np.linalg.inv(eps_conv_matrix)
    mu_conv_inv = np.linalg.inv(mu_conv_matrix)
    p11 = kx_matrix @ eps_conv_inv @ ky_matrix
    p12 = mu_conv_matrix - (kx_matrix @ eps_conv_inv @ kx_matrix)
    p21 = (ky_matrix @ eps_conv_inv @ ky_matrix) - mu_conv_matrix
    p22 = - ky_matrix @ eps_conv_inv @ kx_matrix

    p_matrix = np.block([[p11, p12],
                         [p21, p22]])

    q11 = kx_matrix @ mu_conv_inv @ ky_matrix
    q12 = eps_conv_matrix - (kx_matrix @ mu_conv_inv @ kx_matrix)
    q21 = (ky_matrix @ mu_conv_inv @ ky_matrix) - eps_conv_matrix
    q22 = - ky_matrix @ mu_conv_inv @ kx_matrix

    q_matrix = np.block([[q11, q12],
                         [q21, q22]])

    return p_matrix, q_matrix


def get_eigen(p_matrix, q_matrix):
    """
    Calculate eigen values and eigen vectors

    Args:
        p_matrix (numpy.ndarray):
        q_matrix (numpy.ndarray):

    Returns:
        tuple[scipy.spare.spmatrix, numpy.ndarray]:
            eigen values and eigen vectors, eigen values are arranged in a diagonal matrix
    """
    omega2 = p_matrix @ q_matrix
    lda2s, w_matrix = np.linalg.eig(omega2)
    lda_matrix = sp.sparse.diags(np.sqrt(lda2s, dtype=complex))
    return lda_matrix, w_matrix


def get_v_matrix(q_matrix, w_matrix, lda_matrix):
    """
    Calculate the V matrix form Q, W and lambda
    Args:
        q_matrix (numpy.ndarray):
        w_matrix (numpy.ndarray):
        lda_matrix (scipy.spare.spmatrix):

    Returns:
        numpy.ndarray:
    """
    return q_matrix @ w_matrix @ sp.sparse.linalg.inv(lda_matrix.tocsc())


def get_f_matrix(w_matrix, v_matrix):
    """
    Calculate the field matrix F
    Args:
        w_matrix (numpy.ndarray):
        v_matrix (numpy.ndarray):

    Returns:
        numpy.ndarray:
    """
    return np.block([[w_matrix, w_matrix],
                     [-v_matrix, v_matrix]])


def get_x_matrix(lda_matrix, frequency, thickness, c_const=2.998e5):
    """
    Calculate the propagation matrix X
    Args:
        lda_matrix (scipy.sparse.spmatrix):
        frequency (float):
        thickness (float):
        c_const (float):

    Returns:
        scipy.sparse.spmatrix:
    """
    k0_func = 2 * np.pi * frequency / c_const
    return (-k0_func * thickness * lda_matrix).expm1() + sp.sparse.eye(lda_matrix.shape[0])


def get_cap_a0_b0_matrix(kx_matrix, ky_matrix, mu_inc, mu_trn, kz_inc_matrix, kz_trn_matrix):
    """
    Calculate the initial (capital, to be distinguished from the small b) A and B matrix
    Args:
        kx_matrix (scipy.sparse.spmatrix):
        ky_matrix (scipy.sparse.spmatrix):
        mu_inc (complex):
        mu_trn (complex):
        kz_inc_matrix (scipy.sparse.spmatrix):
        kz_trn_matrix (scipy.sparse.spmatrix):

    Returns:
        tuple[scipy.sparse.spmatrix, scipy.sparse.spmatrix]:
            matrix A and matrix B
    """
    block_size = kx_matrix.shape[0]
    kz_inc_inv_matrix = sp.sparse.diags(1 / kz_inc_matrix.diagonal())
    kz_trn_inv_matrix = sp.sparse.diags(1 / kz_trn_matrix.diagonal())
    identy_matrix = sp.sparse.eye(block_size)
    a31 = -1j / mu_inc * kx_matrix @ ky_matrix @ kz_inc_inv_matrix
    a32 = -1j / mu_inc * (ky_matrix @ ky_matrix + kz_inc_matrix @ kz_inc_matrix) @ kz_inc_inv_matrix
    a41 = 1j / mu_inc * (kx_matrix @ kx_matrix + kz_inc_matrix @ kz_inc_matrix) @ kz_inc_inv_matrix
    a42 = -a31

    b31 = 1j / mu_trn * kx_matrix @ ky_matrix @ kz_trn_inv_matrix
    b32 = 1j / mu_trn * (ky_matrix @ ky_matrix + kz_trn_matrix @ kz_trn_matrix) @ kz_trn_inv_matrix
    b41 = -1j / mu_trn * (kx_matrix @ kx_matrix + kz_trn_matrix @ kz_trn_matrix) @ kz_trn_inv_matrix
    b42 = -b31

    cap_a_matrix = sp.sparse.bmat([[identy_matrix, None],
                                   [None, identy_matrix],
                                   [a31, a32],
                                   [a41, a42]])

    cap_b_matrix = sp.sparse.bmat([[identy_matrix, None],
                                   [None, identy_matrix],
                                   [b31, b32],
                                   [b41, b42]])

    return cap_a_matrix, cap_b_matrix


def get_a_b_matrix(f_matrix, cap_b_matrix):
    """
    Calculate the (small) a and b matrix of certain layer
    Args:
        f_matrix (numpy.ndarray):
        cap_b_matrix (scipy.sparse.spmatrix or numpy.ndarray):

    Returns:
        tuple[numpy.ndarray, numpy.ndarray]:
            a, b matrix
    """
    block_size = f_matrix.shape[0] // 2
    ab = np.linalg.inv(f_matrix) @ cap_b_matrix
    return ab[:block_size, :], ab[block_size:, :]


def get_new_cap_b_matrix(a_matrix, b_matrix, x_matrix, f_matrix):
    """
    Calculate the new B matrix from
    Args:
        a_matrix (numpy.ndarray):
        b_matrix (numpy.ndarray):
        x_matrix (scipy.sparse.spmatrix):
        f_matrix (numpy.ndarray):

    Returns:
        numpy.ndarray
    """
    block_size = x_matrix.shape[0]
    a_inv_matrix = np.linalg.inv(a_matrix)
    identity_matrix = sp.sparse.eye(block_size)
    matrix2 = sp.sparse.bmat([[identity_matrix, None],
                              [None, x_matrix]])
    matrix3 = np.block([[np.eye(block_size)],
                        [b_matrix @ a_inv_matrix @ x_matrix]])
    return f_matrix @ matrix2 @ matrix3


def get_source(px, py, kx_inc, ky_inc, kz_inc_matrix, mu_inc):
    """

    Args:
        px (complex): polarization in x direction
        py (complex): polarization in y direction
        kx_inc (float):
        ky_inc (float):
        kz_inc_matrix (scipy.sparse.spmatrix):
        mu_inc (complex):

    Returns:
        scipy.sparse.spmatrix: the source array (column vector)
    """
    block_size = kz_inc_matrix.shape[0]
    zero_index = (block_size - 1) // 2
    kz_inc = kz_inc_matrix.diagonal()[zero_index]
    pz = - (kx_inc * px + ky_inc * py) / kz_inc
    s1_data = np.array([px])
    s2_data = np.array([py])
    s3_data = np.array([1j / mu_inc * (kz_inc * py - ky_inc * pz)])
    s4_data = np.array([1j / mu_inc * (kx_inc * pz - kz_inc * px)])
    row = np.array([zero_index])
    col = np.array([0])

    s1 = sp.sparse.csc_matrix((s1_data, (row, col)), shape=(block_size, 1))
    s2 = sp.sparse.csc_matrix((s2_data, (row, col)), shape=(block_size, 1))
    s3 = sp.sparse.csc_matrix((s3_data, (row, col)), shape=(block_size, 1))
    s4 = sp.sparse.csc_matrix((s4_data, (row, col)), shape=(block_size, 1))

    return sp.sparse.vstack([s1, s2, s3, s4])


def get_r_t1(cap_a_matrix, cap_b_matrix, source):
    """
    Calculate the r and t1 array
    Args:
        cap_a_matrix (scipy.sparse.spmatrix):
        cap_b_matrix (numpy.ndarray):
        source (scipy.sparse.spmatrix):

    Returns:
        tuple[numpy.ndarray, numpy.ndarray]
    """
    factor = np.hstack((- cap_a_matrix.toarray(), cap_b_matrix))
    rt = np.linalg.inv(factor) @ source
    block_size = rt.shape[0] // 2

    return rt[:block_size, :], rt[block_size:, :]


def get_diff_efficiency(r_vector, t_vector, kx_matrix, ky_matrix, kz_inc_matrix, kz_trn_matrix, mu_inc, mu_trn):
    """

    Args:
        r_vector (numpy.ndarray):
        t_vector (numpy.ndarray):
        kx_matrix (scipy.sparse.spmatrix):
        ky_matrix (scipy.sparse.spmatrix):
        kz_inc_matrix (scipy.sparse.spmatrix):
        kz_trn_matrix (scipy.sparse.spmatrix):
        mu_inc (complex):
        mu_trn (complex):

    Returns:
        tuple[numpy.ndarray, numpy.ndarray]:

    Note: remain to be fixed for lossy inc and trn material
    """
    kz_ref_matrix = -kz_inc_matrix
    block_size = r_vector.shape[0] // 2
    kz_size = kz_inc_matrix.shape[0]
    kz_zero_index = (kz_size - 1) // 2
    kz_inc = kz_inc_matrix.diagonal()[kz_zero_index]
    kz_ref_inv_matrix = sp.sparse.diags(1 / kz_ref_matrix.diagonal())
    kz_trn_inv_matrix = sp.sparse.diags(1 / kz_trn_matrix.diagonal())

    rx_vector = r_vector[:block_size]
    ry_vector = r_vector[block_size:]
    tx_vector = t_vector[:block_size]
    ty_vector = t_vector[block_size:]

    rz_vector = - kz_ref_inv_matrix @ (kx_matrix @ rx_vector + ky_matrix @ ry_vector)
    tz_vector = - kz_trn_inv_matrix @ (kx_matrix @ tx_vector + ky_matrix @ ty_vector)

    r2_vector = np.abs(rx_vector) ** 2 + np.abs(ry_vector) ** 2 + np.abs(rz_vector) ** 2
    t2_vector = np.abs(tx_vector) ** 2 + np.abs(ty_vector) ** 2 + np.abs(tz_vector) ** 2

    reflection_vector = np.real(kz_inc_matrix / mu_inc) / np.real(kz_inc / mu_inc) * r2_vector
    transmission_vector = np.real(kz_trn_matrix / mu_trn) / np.real(kz_inc / mu_inc) * t2_vector

    return reflection_vector, transmission_vector


def get_rcwa_2d(lattice_vector_1, lattice_vector_2, eps_inc, mu_inc, eps_trn, mu_trn, kx_inc, ky_inc, px_inc, py_inc,
                k1_order, k2_order, eps_mesh_list, mu_mesh_list, thickness_list, frequency, unit='nm/THz'):
    """

    Args:
        lattice_vector_1:
        lattice_vector_2:
        eps_inc:
        mu_inc:
        eps_trn:
        mu_trn:
        kx_inc:
        ky_inc:
        px_inc:
        py_inc:
        k1_order:
        k2_order:
        eps_mesh_list:
        mu_mesh_list:
        thickness_list:
        frequency:
        unit:

    Returns:

    """
    if unit == 'nm/THz':
        c_const = 2.998e5
    elif unit == 'mm/GHz':
        c_const = 299.8
    else:
        raise ValueError('unit should be either "nm/THz" or "mm/GHz"')

    # read number of layers
    num_layers = len(eps_mesh_list)

    # initialize reciprocal lattice vector
    reciprocal_vector_1, reciprocal_vector_2 = get_reciprocal(lattice_vector_1, lattice_vector_2)

    # get kx_mat, ky_mat
    kx_matrix, ky_matrix = get_k_matrix(kx_inc, ky_inc, k1_order, k2_order, reciprocal_vector_1, reciprocal_vector_2,
                                        frequency, c_const)

    # get kz_inc_mat, kz_trn_mat
    kz_inc_matrix, kz_trn_matrix = get_kz_matrix(kx_matrix, ky_matrix, eps_inc, mu_inc, eps_trn, mu_trn)

    # get cap_a, cap_b0 matrix
    cap_a_matrix, cap_b_matrix = get_cap_a0_b0_matrix(kx_matrix, ky_matrix, mu_inc, mu_trn, kz_inc_matrix,
                                                      kz_trn_matrix)

    a_matrix_list = []
    x_matrix_list = []

    for i_layer in range(num_layers - 1, -1, -1):
        eps_conv_matrix = conv_matrix.conv_matrix_2d(eps_mesh_list[i_layer], k1_order, k2_order)
        mu_conv_matrix = conv_matrix.conv_matrix_2d(mu_mesh_list[i_layer], k1_order, k2_order)
        thickness = thickness_list[i_layer]

        p_matrix, q_matrix = get_p_q_matrix(kx_matrix, ky_matrix, eps_conv_matrix, mu_conv_matrix)

        lda_matrix, w_matrix = get_eigen(p_matrix, q_matrix)

        v_matrix = get_v_matrix(q_matrix, w_matrix, lda_matrix)

        f_matrix = get_f_matrix(w_matrix, v_matrix)

        x_matrix = get_x_matrix(lda_matrix, frequency, thickness, c_const)

        a_matrix, b_matrix = get_a_b_matrix(f_matrix, cap_b_matrix)

        new_cap_b_matrix = get_new_cap_b_matrix(a_matrix, b_matrix, x_matrix, f_matrix)

        cap_b_matrix = new_cap_b_matrix

        a_matrix_list.append(a_matrix)
        x_matrix_list.append(x_matrix)

    a_matrix_list.reverse()
    x_matrix_list.reverse()

    # calculate source
    s_vector = get_source(px_inc, py_inc, kx_inc, ky_inc, kz_inc_matrix, mu_inc)

    # calculate r, t1
    r_vector, t_vector = get_r_t1(cap_a_matrix, cap_b_matrix, s_vector)

    # loop over layers
    for i_layer in range(num_layers):
        a_matrix = a_matrix_list[i_layer]
        x_matrix = x_matrix_list[i_layer]
        t_vector = np.linalg.inv(a_matrix) @ x_matrix @ t_vector

    ref_vector, trn_vector = get_diff_efficiency(r_vector, t_vector, kx_matrix, ky_matrix, kz_inc_matrix, kz_trn_matrix,
                                                 mu_inc, mu_trn)
    return ref_vector, trn_vector


if __name__ == "__main__":
    """ the code is tested for diffraction grating system of https://www.comsol.com/model/diffraction-grating-19083 """
    import matplotlib.pyplot as plt

    # %% main loop

    """ input start """
    # initialize
    eps_list = []
    mu_list = []
    thick_list = []

    # input layers of materials
    N_layers = 1  # number of layers
    n_sio2 = 1.5487
    eps_sio2 = n_sio2 ** 2

    d = 340
    lattice_vec_1 = np.array([d, 0])  # lattice vector 1
    lattice_vec_2 = np.array([0, d])  # lattice vector 2

    eps_i = 1  # eps and mu of incident material
    mu_i = 1

    eps_t = eps_sio2  # eps and mu of transmitted material
    mu_t = 1

    eps_list.append(np.hstack((eps_sio2 * np.ones((2, 256)), np.ones((2, 256)))))
    # eps distribution of layer 1. use 2 rows to be compatible with FFT
    mu_list.append(np.ones((2, 512)))  # mu distribution of layer 1
    thick_list.append(d / 4)

    # input frequency, kx_inc, ky_inc, k1_order, k2_order
    freq = 679.8
    c = 2.998e5
    k0 = 2 * np.pi * freq / c

    ref_0_order_list = []
    ref_p1_order_list = []
    ref_m1_order_list = []

    trn_0_order_list = []
    trn_p1_order_list = []
    trn_m1_order_list = []

    angle_array = np.linspace(0, 89.9 / 180 * np.pi, 90)
    for angle_i in angle_array:
        kx_i = k0 * np.sin(angle_i)
        ky_i = 0

        k1_ord = 40
        k2_ord = 0

        # input source polarization
        px_s = 0
        py_s = 1

        """ input end """

        """ calculation start """

        ref_vec, trn_vec = get_rcwa_2d(lattice_vec_1, lattice_vec_2, eps_i, mu_i, eps_t, mu_t, kx_i, ky_i, px_s, py_s,
                                       k1_ord, k2_ord, eps_list, mu_list, thick_list, freq)
        order_0_index = (ref_vec.shape[0] - 1) // 2

        ref_0_order_list.append(ref_vec[order_0_index, 0])
        ref_p1_order_list.append(ref_vec[order_0_index - 1, 0])  # the index is negative
        ref_m1_order_list.append(ref_vec[order_0_index + 1, 0])  # because the index is counted from the largest order

        trn_0_order_list.append(trn_vec[order_0_index, 0])
        trn_p1_order_list.append(trn_vec[order_0_index - 1, 0])
        trn_m1_order_list.append(trn_vec[order_0_index + 1, 0])

    diff_eff = np.vstack((angle_array / np.pi * 180,
                          np.array(ref_m1_order_list),
                          np.array(ref_0_order_list),
                          np.array(ref_p1_order_list),
                          np.array(trn_m1_order_list),
                          np.array(trn_0_order_list),
                          np.array(trn_p1_order_list)))
    diff_eff = diff_eff.transpose()
    np.savetxt('diff_eff.csv', diff_eff, delimiter=',',
               fmt='%.3f', header='angle(deg),ref_m1,ref_0,ref_p1,trn_m1,trn_0,trn_p1')

    plt.figure(1, figsize=(4, 3), dpi=200)
    plt.plot(angle_array / np.pi * 180, ref_0_order_list, label='r0')
    plt.plot(angle_array / np.pi * 180, ref_m1_order_list, label='r-1')
    plt.plot(angle_array / np.pi * 180, ref_p1_order_list, label='r+1')
    plt.xlabel('angle (deg)')
    plt.ylabel('Reflection')
    plt.legend()

    plt.figure(2, figsize=(4, 3), dpi=200)
    plt.plot(angle_array / np.pi * 180, trn_0_order_list, label='t0')
    plt.plot(angle_array / np.pi * 180, trn_m1_order_list, label='t-1')
    plt.plot(angle_array / np.pi * 180, trn_p1_order_list, label='t+1')
    plt.legend()
    plt.xlabel('angle (deg)')
    plt.ylabel('Transmission')

    plt.show()
